Clamp slice bounds when a slice is combined with an array index - #4513
Open
IshaanPotle wants to merge 1 commit into
Open
IshaanPotle wants to merge 1 commit into
IshaanPotle wants to merge 1 commit into
Conversation
When a slice is combined with an array index, the slice bounds are
resolved for negative values but never clamped into [0, axis_size], so
out-of-range bounds are passed to arange() verbatim:
start = (start < 0) ? start + src.shape(i) : start;
A single adjustment is not enough: for a 4-element axis, -100 becomes
-96, and arange(-96, 100) yields 196 indices into that axis.
Consequences on a (4, 5) array with idx = [0, 1]:
a[0:100, idx] -> (100, 2) of repeated data, NumPy gives (4, 2)
a[-100:4, idx] -> (100, 2), NumPy gives (4, 2)
a[-100:-50, idx] -> (50, 2), NumPy gives (0, 2)
a[-100:4, idx] = 0 -> writes 1 row, NumPy writes 4
a[100:-100:-1, idx] -> SIGBUS, NumPy gives (4, 2)
The last case is an out-of-bounds read reachable from ordinary Python.
Slicing past the end is idiomatic and safe elsewhere in Python, so the
first case in particular is easy to hit by accident, and none of these
raise.
Add adjust_slice_bounds(), which mirrors CPython's PySlice_AdjustIndices
(clamping to [0, n] for positive strides and [-1, n-1] for negative
ones), and use it in both mlx_gather_nd and mlx_scatter_args_nd, which
carried the same partial adjustment.
Plain slicing and slice-plus-integer indexing were already correct and
are untouched.
Tested with a differential sweep against NumPy over 11 x 11 x 7
start/stop/step combinations for both __getitem__ and __setitem__.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
When combining a slice with an array of indices (advanced indexing), out-of-range slice bounds are adjusted (start + axis_size), but never clamped to the actual dimension size as standard Python/NumPy semantics require. Because these bounds remain unclamped, index calculations generate invalid memory offsets:
Positive steps silently return incorrect shapes or corrupt row data.
Negative steps can access unmapped memory, resulting in a process crash (SIGBUS).
Note: Plain slicing (a[0:100]) and slice + integer indexing (a[0:100, 1]) were already correct and remain completely untouched by this change.
Minimal Reproducible Example
Proposed Fix
Added a clamping helper mirroring PySlice_AdjustIndices to both the read and write indexing paths within the Python bindings. This ensures slice bounds are constrained to valid ranges prior to index offset calculation, matching NumPy's behavior for these advanced indexing patterns.
Verification & Test Coverage
Differential Testing: Verified against NumPy across 1,694 differential assertions (sweeping a 2D array and index array across 847 parameter combinations). Results match NumPy for these cases.
Local Test Suite: Ran 7 CPU test files (test_array, test_ops, test_autograd, test_vmap, test_compile, test_einsum, test_blas) with -DMLX_BUILD_METAL=OFF (434 passed, 26 skipped). Because this fix resides strictly in device-independent index arithmetic within the Python bindings, CPU coverage validates the logical correctness; Metal/GPU execution paths will be exercised in CI.
Issue Tracking: Closes #4399 (additionally resolves an unreported SIGBUS crash during negative-step slicing combined with array indexing).